package com.xiejun.aio;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousChannelGroup;
import java.nio.channels.AsynchronousServerSocketChannel;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class AioServer {
    private final int port;
    private AsynchronousServerSocketChannel serverSocketChannel;
    private AsynchronousChannelGroup channelGroup;

    public AioServer(int port) {
        this.port = port;
    }

    public void start() throws IOException {
        // 创建线程池
        ExecutorService executorService = Executors.newFixedThreadPool(20);
        
        // 创建通道组
        channelGroup = AsynchronousChannelGroup.withThreadPool(executorService);
        
        // 打开服务端通道
        serverSocketChannel = AsynchronousServerSocketChannel.open(channelGroup);
        
        // 绑定端口
        serverSocketChannel.bind(new InetSocketAddress(port));
        
        System.out.println("AIO Server started on port " + port);
        
        // 监听客户端连接
        serverSocketChannel.accept(this, new AcceptHandler());
        
        // 保持服务器运行
        try {
            Thread.sleep(Integer.MAX_VALUE);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    private static class AcceptHandler implements CompletionHandler<AsynchronousSocketChannel, AioServer> {
        @Override
        public void completed(AsynchronousSocketChannel channel, AioServer attachment) {
            // 继续监听下一个连接
            attachment.serverSocketChannel.accept(attachment, this);
            
            // 处理当前连接
            System.out.println("New client connected: " + channel);
            ByteBuffer buffer = ByteBuffer.allocate(1024);
            channel.read(buffer, buffer, new ReadHandler(channel));
        }

        @Override
        public void failed(Throwable exc, AioServer attachment) {
            exc.printStackTrace();
        }
    }

    private static class ReadHandler implements CompletionHandler<Integer, ByteBuffer> {
        private final AsynchronousSocketChannel channel;

        public ReadHandler(AsynchronousSocketChannel channel) {
            this.channel = channel;
        }

        @Override
        public void completed(Integer result, ByteBuffer attachment) {
            if (result == -1) {
                try {
                    channel.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
                return;
            }

            attachment.flip();
            byte[] bytes = new byte[attachment.remaining()];
            attachment.get(bytes);
            String message = new String(bytes);
            System.out.println("Received message: " + message);

            // 回复客户端
            String reply = "Echo: " + message;
            ByteBuffer buffer = ByteBuffer.wrap(reply.getBytes());
            channel.write(buffer, buffer, new WriteHandler(channel));
        }

        @Override
        public void failed(Throwable exc, ByteBuffer attachment) {
            exc.printStackTrace();
        }
    }

    private static class WriteHandler implements CompletionHandler<Integer, ByteBuffer> {
        private final AsynchronousSocketChannel channel;

        public WriteHandler(AsynchronousSocketChannel channel) {
            this.channel = channel;
        }

        @Override
        public void completed(Integer result, ByteBuffer attachment) {
            if (attachment.hasRemaining()) {
                channel.write(attachment, attachment, this);
            } else {
                // 继续读取客户端数据
                attachment.clear();
                channel.read(attachment, attachment, new ReadHandler(channel));
            }
        }

        @Override
        public void failed(Throwable exc, ByteBuffer attachment) {
            exc.printStackTrace();
        }
    }

    public static void main(String[] args) {
        try {
            new AioServer(9001).start();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}